#include "bi_reconstruction.hpp"
#include "bi_constants.hpp"
#include "reconstruction.hpp"
#include "reconstructionHelpers.hpp"
#include "solverBase.hpp"
#include "sweSolver.hpp"
#include <catch2/catch_test_macros.hpp>

using namespace hms;
using namespace hms::swe;

TEST_CASE( "Hydrostatic reconstruction" ){
	SECTION("Blockwise"){
		constexpr Index nEdges {3};

		ArrayXXs buf (nEdges, 2*nReconstComp);
		EdgesValues ev { mapEdgesValues(buf, nEdges) };

		Array1XsMap&
			dL {ev.dLR[0]}, dR {ev.dLR[1]},
			zL {ev.zLR[0]}, zR {ev.zLR[1]},
			hL {ev.hLR[0]}, hR {ev.hLR[1]};
		// Array2XsMap&
		// 	vL {ev.vLR[0]}, vR {ev.vLR[1]};

		zL << 0, 0, 1.1;
		zR << 0, 1, 0;
		hL << 1, 1, 1;
		hR << 1, 1, 1.5;
		reconstruction::hydrostatic_firstOrder(ev);
		REQUIRE( (zL == zR).all() );
		REQUIRE( dL.isApprox(RowArray3s{1, 0, 0}) );
		REQUIRE( dR.isApprox(RowArray3s{1, 0, 0.4}) );
	}
}

TEST_CASE( "Edges values reconstruction" ){
	using SolverType = SWESolver<UniMesh>;

	constexpr Index nEdges {3};
	ArrayXXs buf (nEdges, 18);

	Pair<RowArray3s> dCellLR, zCellLR, hCellLR, dCellLLRR, zCellLLRR, hCellLLRR;
	Pair<ArrayNNs<2,3>> vCellLR, qCellLR, vCellLLRR, qCellLLRR;
	dCellLR[0] << 0, 1, 0.9*SolverType::minWaterDepth;
	dCellLR[1] << 1, 0, 0;
	zCellLR[0] = 0;
	zCellLR[1] = 0;
	vCellLR[0] = 1;
	vCellLR[1] = 1;
	dCellLLRR[0] << 1, 0, 0;
	dCellLLRR[1] << 0, 1, 0.9*SolverType::minWaterDepth;
	zCellLLRR[0] = 0;
	zCellLLRR[1] = 0;
	for (int i{0}; i<2; ++i){
		hCellLR[i] = zCellLR[i] + dCellLR[i];
		hCellLLRR[i] = zCellLLRR[i] + dCellLLRR[i];
		qCellLR[i] = (dCellLR[i] < SolverType::minWaterDepth).replicate(2,1)
			.select( 0, dCellLR[i].replicate(2,1) * vCellLR[i] );
		qCellLLRR[i] = (dCellLLRR[i] < SolverType::minWaterDepth).replicate(2,1)
			.select( 0, dCellLLRR[i].replicate(2,1) * vCellLLRR[i] );
	}
	EdgesValues ev { mapEdgesValues( buf, dCellLR[0].size() ) };

	SECTION("First order"){
		reconstruction::firstOrder<SolverType>(ev, dCellLR, zCellLR, vCellLR);

		REQUIRE( ev.dLR[0].isApprox(RowArray3s{0, 1, 0}) );
		REQUIRE( ev.dLR[1].isApprox(RowArray3s{1, 0, 0}) );
		REQUIRE( (ev.hLR[0] == ev.dLR[0]).all() );
		REQUIRE( (ev.hLR[1] == ev.dLR[1]).all() );
		REQUIRE( (ev.vLR[0].reshaped() == ArrayNs<6>{0, 0, 1, 1, 0, 0}).all() );
		REQUIRE( (ev.vLR[1].reshaped() == ArrayNs<6>{1, 1, 0, 0, 0, 0}).all() );
	}
	SECTION("Central"){
		reconstruction::central<SolverType>(ev, dCellLR, hCellLR, qCellLR);

		REQUIRE( ev.dLR[0].isApprox(RowArray3s{0.5, 0.5, 0}) );
		REQUIRE( (ev.dLR[1] == ev.dLR[0]).all() );
		REQUIRE( ev.hLR[0].isApprox(RowArray3s{0.5, 0.5, 0}) );
		REQUIRE( (ev.hLR[1] == ev.hLR[0]).all() );
		REQUIRE( ev.zLR[0].isApproxToConstant(0) );
		REQUIRE( ev.zLR[1].isApproxToConstant(0) );
		REQUIRE( ev.vLR[0].reshaped().isApprox(ArrayNs<6>{1, 1, 1, 1, 0, 0}) );
		REQUIRE( (ev.vLR[1] == ev.vLR[0]).all() );
	}
	SECTION("Higher order"){
		auto lim { LimiterFunction::minMod };
		reconstruction::highOrder<SolverType>(
			ev,
			dCellLR, dCellLLRR,
			hCellLR, hCellLLRR,
			qCellLR, qCellLLRR,
			1,1,
			lim
		);
		REQUIRE( ev.dLR[0].isApprox(RowArray3s{0, 1, 0}) );
		REQUIRE( ev.dLR[1].isApprox(RowArray3s{1, 0, 0}) );
		REQUIRE( ev.hLR[0].isApprox(RowArray3s{0, 1, 0}) );
		REQUIRE( ev.hLR[1].isApprox(RowArray3s{1, 0, 0}) );
		REQUIRE( ev.vLR[0].reshaped().isApprox(ArrayNs<6>{0, 0, 1, 1, 0, 0}) );
		REQUIRE( ev.vLR[1].reshaped().isApprox(ArrayNs<6>{1, 1, 0, 0, 0, 0}) );
	}
}

TEST_CASE("Blockwise MUSCL scheme"){
	/* we want to reconstruct edge values with a TVD scheme. That's why we need
	 * four cells per edge: to be able to calculate slopes */
	using reconstruction::musclScheme;
	using reconstruction::mapPair;

	Array2Xs vLR_buf, vCellLR_buf, vCellLLRR_buf;
	vLR_buf.resize(2,4);
	vCellLR_buf.resizeLike(vLR_buf);
	vCellLLRR_buf.resizeLike(vLR_buf);

	Array2XsMap
		vCellL  {vCellLR_buf  .col(0).data(), 2, 2},
		vCellR  {vCellLR_buf  .col(2).data(), 2, 2},
		vCellLL {vCellLLRR_buf.col(0).data(), 2, 2},
		vCellRR {vCellLLRR_buf.col(2).data(), 2, 2};

	/* to test out the function, we want three different cases:
	 * r < 0, 0 < r < 1 and 1 < r, with rL = (vL - vLL)/(vR-vL).
	 * Therefore, we need the cases
	 * 1. vL < vLL && vR > vL || vL > vLL && vR < vL,
	 * 2. vL - vLL < vR - vL (either both larger or both smaller than zero) and
	 * 3. vL - vLL > vR - vL
	 * As a fourth case, vL == vR provokes division by zero.
	 *  */
	vCellL <<
		 0, // 1.
		 1, // 2.
		 2, // 3.
		-1; // 4.
	vCellLL <<
		 1, // 1.
		 0, // 2.
		 0, // 3.
		 0;
	vCellR <<
		 1, // 1.
		 3, // 2.
		 3, // 3.
		-1; // 4.
	vCellRR <<
		 1, // 2., both negative
		 1, // 1.
		 5, // 3., both negative
		 0;
	Pair<Array2XsMap>
		vCellLR   {{ std::move(vCellL ), std::move(vCellR ) }},
		vCellLLRR {{ std::move(vCellLL), std::move(vCellRR) }};
	Pair<Array2XsMap> vLR { mapPair<2>(vLR_buf, 2) };

	Array2Xs limFuncBuf (2,4); // as many entries as vCellLR
	auto lim { LimiterFunction::minMod };

	auto minMod = [](scalar r) -> scalar {
		return std::max(0., std::min(1., r));
	};
	Pair<Array2XsMap> psiMap { mapPair<2>(limFuncBuf, 2) };

	/* Because the initialisation with operator<< is RowMajor, but the
	 * resulting matrices are ColMajor, the second and third entry are swapped.
	 */
	SECTION("UniMesh"){
		musclScheme(vLR, psiMap, vCellLR, vCellLLRR, 1, 1, lim);
		auto vL_re { vLR[0].reshaped() };
		/* phiL + 0.5 * psi( (phiL - phiLL)/(phiR - phiL) ) * (phiR-phiL) */
		REQUIRE( vLR[0].reshaped().isApprox(Array4s{
			 0 + 0.5 * minMod(-1./1) * 1,
			 2 + 0.5 * minMod( 2./1) * 1,
			 1 + 0.5 * minMod( 1./2) * 2,
			-1 + 0.5 * minMod( 0   ) * 0
		}) );
		/* phiR + 0.5 * psi( (phiR - phiRR)/(phiL - phiR) ) * (phiL-phiR) */
		REQUIRE( vLR[1].reshaped().isApprox(Array4s{
			 1 + 0.5 * minMod( 0   ) * (-1),
			 3 + 0.5 * minMod( 2./1) * (-1),
			 3 + 0.5 * minMod(-2./2) * (-2),
			-1 + 0.5 * minMod( 0   ) * 0
		}) );
	}
	SECTION("RectMesh"){
		using ER = EdgeRange;
		Array4s dx { 1, 2, 2, 4 };
		RowArrayXs
			dxL { dxQuotient<RectMesh, ER::vertical, 0>(dx) },
			dxR { dxQuotient<RectMesh, ER::vertical, 1>(dx) };

		musclScheme(vLR, psiMap, vCellLR, vCellLLRR, dxL, dxR, lim);

		auto vL_re { vLR[0].reshaped() };
		/* phiL + 0.5 * psi( (phiL - phiLL)/(phiR - phiL) * dx/dxL ) * (phiR-phiL) */
		REQUIRE( vLR[0].reshaped().isApprox(Array4s{
			 0 + 0.5 * minMod(-1./1 * 2) * 1,
			 2 + 0.5 * minMod( 2./1 * 2) * 1,
			 1 + 0.5 * minMod( 1./2 * 1) * 2,
			-1 + 0.5 * minMod( 0       ) * 0
		}) );
		/* phiR + 0.5 * psi( (phiR - phiRR)/(phiL - phiR) * dx/dxR ) * (phiL-phiR) */
		REQUIRE( vLR[1].reshaped().isApprox(Array4s{
			 1 + 0.5 * minMod( 0         ) * (-1),
			 3 + 0.5 * minMod( 2./1 * 1  ) * (-1),
			 3 + 0.5 * minMod(-2./2 * 0.5) * (-2),
			-1 + 0.5 * minMod( 0         ) * 0
		}) );
	}
}

TEST_CASE("Slope limiter"){
	using reconstruction::setSlopeLimiter;
	SECTION("Blockwise"){
		ArrayNNs<2,2> r, psi;
		r.col(0) << -1, 0.5; // left
		r.col(1) <<  1, 2; // right
		using LF = LimiterFunction;
		LF lim;
		SECTION("Limiter: none"){
			lim = LF::none;
			setSlopeLimiter(psi, r, lim);
			REQUIRE( (psi.col(0) == 1).all() );
			REQUIRE( (psi.col(1) == 1).all() );
		}
		SECTION("Limiter: minMod"){
			lim = LF::minMod;
			setSlopeLimiter(psi, r, lim);
			REQUIRE( psi.col(0).isApprox(Array2s{0, 0.5}) );
			REQUIRE( psi.col(1).isApprox(Array2s{1, 1}) );
		}
		SECTION("Limiter: vanAlbada"){
			lim = LF::vanAlbada;
			setSlopeLimiter(psi, r, lim);
			REQUIRE( psi.col(0).isApprox(Array2s{0, 0.75/1.25}) );
			REQUIRE( psi.col(1).isApprox(Array2s{1, 1.2}) );
		}
	}
}
